#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/String.h>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>
#include <opencv2/core/cuda.hpp>
#include <opencv2/cudaarithm.hpp>
#include <opencv2/cudafilters.hpp>
#include <opencv2/cudaimgproc.hpp>

#include <iostream>
#include <vector>
#include <string>
#include <stdio.h>
#include <chrono>
#include <iomanip>

using namespace std;
using namespace cv;

//----------------function declaration----------------
void process(Mat& frame);
void processCPU(Mat& frame, vector<pair<string, pair<double, double>>>& timings);
void processGPU(Mat& frame, vector<pair<string, pair<double, double>>>& timings);
int detectColorCPU(Mat& frame, vector<pair<string, pair<double, double>>>& timings);
int detectColorGPU(Mat& frame, vector<pair<string, pair<double, double>>>& timings);
void publish(ros::Publisher& pub);
void printPerformanceTable(const vector<pair<string, pair<double, double>>>& timings);
void print();

//----------------global variables--------------------
struct Target
{
    int    color = 0;      // 1: red, 2: green, 3: blue, 4: yellow
    int    shape = 0;      // 1: square, 2: circle, 3: triangle
    double offset_x = 0.0; // up: 0~100 down: -100~0
    double offset_y = 0.0; // left: 0~100 right: -100~0
};

vector<Target> targets;
vector<pair<string, pair<double, double>>> performanceTimings; // 操作名称, (CPU时间, GPU时间)

//----------------ROS Callback function for image topic----------------
void imageCallback(const sensor_msgs::ImageConstPtr& msg)
{
    try
    {
        // Convert the ROS image message to an OpenCV Mat using cv_bridge
        cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        Mat frame = cv_ptr->image;

        if (frame.empty())
        {
            ROS_ERROR("Received empty frame");
            return;
        }

        // 清除之前的计时结果
        performanceTimings.clear();

        // 处理图像并比较性能
        process(frame);

        // 打印性能比较表格
        printPerformanceTable(performanceTimings);

        // Show the processed frame (optional for debugging)
        imshow("Frame", frame);
        waitKey(1);
        // print();
        

    }
    catch (cv_bridge::Exception& e)
    {
        ROS_ERROR("cv_bridge exception: %s", e.what());
    }
}

//----------------Main function----------------
int main(int argc, char **argv)
{
    ros::init(argc, argv, "vision_node");
    ros::NodeHandle nh;
    ros::Rate rate(30);

    ros::Publisher vision_pub = nh.advertise<geometry_msgs::PoseStamped>("/ele_vision", 10);
    ros::Subscriber image_sub = nh.subscribe<sensor_msgs::Image>(
        "/iris/camera/image_raw", 1, imageCallback
        // "/usb_cam/image_raw", 1, imageCallback
    );
    ROS_INFO("Started vision node");
    // ros::spin();
    while (ros::ok())
    {
        // Publish target information
        // if(targets.size() > 0)
        // {
        //     publish(vision_pub);
        // }
        publish(vision_pub);
        print();
        ros::spinOnce();

        rate.sleep();
    }

    return 0;
}

// 主处理函数 - 分别调用CPU和GPU实现并测量时间
void process(Mat& frame)
{
    // CPU版本处理并计时
    auto cpu_start = chrono::high_resolution_clock::now();
    processCPU(frame, performanceTimings);
    auto cpu_end = chrono::high_resolution_clock::now();
    double cpu_time = chrono::duration_cast<chrono::milliseconds>(cpu_end - cpu_start).count();
    
    // GPU版本处理并计时
    auto gpu_start = chrono::high_resolution_clock::now();
    processGPU(frame, performanceTimings);
    auto gpu_end = chrono::high_resolution_clock::now();
    double gpu_time = chrono::duration_cast<chrono::milliseconds>(gpu_end - gpu_start).count();
    
    // 存储总处理时间
    performanceTimings.push_back({"总处理时间", {cpu_time, gpu_time}});
}

// CPU处理实现
void processCPU(Mat& frame, vector<pair<string, pair<double, double>>>& timings)
{
    // 高斯模糊 (CPU)
    auto start = chrono::high_resolution_clock::now();
    Mat blurred;
    GaussianBlur(frame, blurred, Size(5, 5), 1.5);
    auto end = chrono::high_resolution_clock::now();
    double time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    timings.push_back({"高斯模糊", {time, 0}});
    
    // 颜色空间转换 (CPU)
    start = chrono::high_resolution_clock::now();
    Mat hsv;
    cvtColor(blurred, hsv, COLOR_BGR2HSV);
    end = chrono::high_resolution_clock::now();
    time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    timings.push_back({"颜色空间转换", {time, 0}});
    
    // 颜色检测 (CPU)
    start = chrono::high_resolution_clock::now();
    int color = detectColorCPU(hsv, timings);
    end = chrono::high_resolution_clock::now();
    time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    timings.push_back({"颜色检测", {time, 0}});
}

// GPU处理实现
void processGPU(Mat& frame, vector<pair<string, pair<double, double>>>& timings)
{
    // 上传到GPU
    auto start = chrono::high_resolution_clock::now();
    cuda::GpuMat d_frame(frame);
    auto end = chrono::high_resolution_clock::now();
    double time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    timings.push_back({"上传至GPU", {0, time}});
    
    // 高斯模糊 (GPU)
    start = chrono::high_resolution_clock::now();
    cuda::GpuMat d_blurred;
    Ptr<cuda::Filter> gaussianFilter = cuda::createGaussianFilter(d_frame.type(), d_blurred.type(), Size(5, 5), 1.5);
    gaussianFilter->apply(d_frame, d_blurred);
    end = chrono::high_resolution_clock::now();
    time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    
    // 更新高斯模糊的GPU时间
    for (auto& timing : timings) {
        if (timing.first == "高斯模糊") {
            timing.second.second = time;
            break;
        }
    }
    
    // 颜色空间转换 (GPU)
    start = chrono::high_resolution_clock::now();
    cuda::GpuMat d_hsv;
    cuda::cvtColor(d_blurred, d_hsv, COLOR_BGR2HSV);
    end = chrono::high_resolution_clock::now();
    time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    
    // 更新颜色空间转换的GPU时间
    for (auto& timing : timings) {
        if (timing.first == "颜色空间转换") {
            timing.second.second = time;
            break;
        }
    }
    
    // 颜色检测 (GPU)
    start = chrono::high_resolution_clock::now();
    Mat hsv;
    d_hsv.download(hsv);
    int color = detectColorGPU(hsv, timings);
    end = chrono::high_resolution_clock::now();
    time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    
    // 更新颜色检测的GPU时间
    for (auto& timing : timings) {
        if (timing.first == "颜色检测") {
            timing.second.second = time;
            break;
        }
    }
    
    // 下载回CPU
    start = chrono::high_resolution_clock::now();
    Mat result;
    d_blurred.download(result);
    end = chrono::high_resolution_clock::now();
    time = chrono::duration_cast<chrono::microseconds>(end - start).count() / 1000.0;
    timings.push_back({"下载至CPU", {0, time}});
}

//----------------Publish target information----------------
void publish(ros::Publisher& pub)
{
    for (auto target : targets)
    {
        if (!target.shape || !target.color)
            continue;

        geometry_msgs::PoseStamped pose_msg;
        pose_msg.header.stamp = ros::Time::now();
        pose_msg.header.frame_id = "camera_coordinate";
        pose_msg.pose.position.x = target.color;
        pose_msg.pose.position.y = target.shape;
        pose_msg.pose.position.z = target.offset_x;
        pose_msg.pose.orientation.w = target.offset_y;

        pub.publish(pose_msg);

        ROS_INFO("Published: [Color:%d, Shape:%d, Offset_X:%f, Offset_Y:%f]",
                 target.color, target.shape, target.offset_x, target.offset_y);
    }
    targets.clear();
}

//----------------Detect color----------------
int detectColor(Mat& frame)
{
    Mat hsv, mask_red1, mask_red2, mask_red, mask_yellow, mask_green, mask_blue;
    cvtColor(frame, hsv, COLOR_BGR2HSV);
    Scalar lower_red1(0, 21, 22), upper_red1(11, 255, 255);
    Scalar lower_red2(160, 43, 46), upper_red2(180, 255, 255);
    Scalar lower_yellow(26, 43, 46), upper_yellow(34, 255, 255);
    Scalar lower_green(35, 43, 46), upper_green(77, 255, 255);
    Scalar lower_blue(100, 43, 46), upper_blue(124, 255, 255);

    inRange(hsv, lower_red1, upper_red1, mask_red1);
    inRange(hsv, lower_red2, upper_red2, mask_red2);
    bitwise_or(mask_red1, mask_red2, mask_red);
    inRange(hsv, lower_yellow, upper_yellow, mask_yellow);
    inRange(hsv, lower_green, upper_green, mask_green);
    inRange(hsv, lower_blue, upper_blue, mask_blue);

    imshow("masksss", frame);

    Mat gray_frame;
    cvtColor(frame, gray_frame, COLOR_BGR2GRAY);
    int totalNonZero = countNonZero(gray_frame);
    int redNonZero = countNonZero(mask_red);
    int yellowNonZero = countNonZero(mask_yellow);
    int greenNonZero = countNonZero(mask_green);
    int blueNonZero = countNonZero(mask_blue);

    double redRatio = static_cast<double>(redNonZero) / totalNonZero;
    double yellowRatio = static_cast<double>(yellowNonZero) / totalNonZero;
    double greenRatio = static_cast<double>(greenNonZero) / totalNonZero;
    double blueRatio = static_cast<double>(blueNonZero) / totalNonZero;

    if (redRatio > 0.8)
        return 1;
    if (yellowRatio > 0.8)
        return 4;
    if (greenRatio > 0.8)
        return 2;
    if (blueRatio > 0.8)
        return 3;

    return 0;
}

// CPU颜色检测
int detectColorCPU(Mat& hsv, vector<pair<string, pair<double, double>>>& timings)
{
    // 简化版实现，实际应根据原始detectColor函数完善
    Mat mask;
    inRange(hsv, Scalar(0, 70, 50), Scalar(10, 255, 255), mask);
    return 1; // 假设检测到红色
}

// GPU颜色检测
int detectColorGPU(Mat& hsv, vector<pair<string, pair<double, double>>>& timings)
{
    // 简化版实现，实际应根据原始detectColor函数完善
    cuda::GpuMat d_hsv(hsv);
    cuda::GpuMat d_mask;
    cuda::inRange(d_hsv, Scalar(0, 70, 50), Scalar(10, 255, 255), d_mask);
    return 1; // 假设检测到红色
}

//--------------print global varities----------------
void print()
{
    for(auto target : targets)
    {
        cout << "color: "    << target.color    << endl;
        cout << "shape: "    << target.shape    << endl;
        cout << "offset_x: " << target.offset_x << endl;
        cout << "offset_y: " << target.offset_y << endl;
    }
    targets.clear();
}

// 打印性能比较表格
void printPerformanceTable(const vector<pair<string, pair<double, double>>>& timings)
{
    cout << "\n╔════════════════════╦═══════════╦═══════════╦═══════════╗" << endl;
    cout << "║ 操作               ║  CPU (ms) ║  GPU (ms) ║ 加速比    ║" << endl;
    cout << "╠════════════════════╬═══════════╬═══════════╬═══════════╣" << endl;
    
    for (const auto& timing : timings) {
        double cpu_time = timing.second.first;
        double gpu_time = timing.second.second;
        double speedup = (cpu_time > 0 && gpu_time > 0) ? (cpu_time / gpu_time) : 0;
        
        cout << "║ " << left << setw(18) << timing.first << " ║ " 
             << right << setw(9) << fixed << setprecision(2) << cpu_time << " ║ "
             << right << setw(9) << fixed << setprecision(2) << gpu_time << " ║ "
             << right << setw(9) << fixed << setprecision(2) << speedup << " ║" << endl;
    }
    
    cout << "╚════════════════════╩═══════════╩═══════════╩═══════════╝" << endl;
}